#!/usr/bin/env python3
"""
简化的 Register API 使用示例
这是一个最简化的示例，展示如何快速使用 register 接口
"""

import requests
import json
import time


def simple_register_example():
    """简化的注册示例"""
    
    # 控制器地址
    controller_url = "http://localhost:57675/controller/register"
    
    # 注册数据 - 修复数据不匹配问题
    register_data = {
        "job_name": "simple_job",
        "model_name": "model_0",
        "role": "worker",
        "pod_ip": "127.0.0.1",
        "host_ip": "192.168.1.100",
        "bussiness_port": ["8080"],  # 只使用一个端口，匹配设备数量
        "mgmt_port": "9090",
        "parallel_config": {
            "dp_size": 1,
            "cp_size": 1,
            "tp_size": 1,
            "sp_size": 1,
            "ep_size": 1,
            "pp_size": 1,
            "world_size": 1
        },
        "ranktable": {
            "version": "1.0",
            "status": "active",
            "server_count": "1",
            "server_list": [
                {
                    "server_id": "server_0",
                    "host_ip": "192.168.1.100",
                    "device": [
                        {
                            "device_id": "0",
                            "device_ip": "192.168.1.100",
                            "super_device_id": None,
                            "rank_id": "0"
                        }
                    ]
                }
            ]
        }
    }
    
    try:
        print("发送注册请求...")
        print(f"数据: {json.dumps(register_data, indent=2, ensure_ascii=False)}")
        
        # 发送请求
        response = requests.post(controller_url, json=register_data, timeout=30)
        
        # 检查响应
        if response.status_code == 200:
            result = response.json()
            print(f"注册成功! 响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
        else:
            print(f"注册失败! 状态码: {response.status_code}")
            print(f"错误信息: {response.text}")
            
    except requests.exceptions.ConnectionError:
        print("连接失败! 请确保控制器服务正在运行")
    except requests.exceptions.Timeout:
        print("请求超时!")
    except Exception as e:
        print(f"发生错误: {e}")


def multi_device_register_example():
    """多设备注册示例"""
    
    controller_url = "http://localhost:57675/controller/register"
    
    # 多设备配置示例
    register_data = {
        "job_name": "multi_device_job",
        "model_name": "model_1",
        "role": "worker",
        "pod_ip": "127.0.0.1",
        "host_ip": "192.168.1.100",
        "bussiness_port": [
            "8080",
            "8081"
        ],
        "mgmt_port": "9090",
        "parallel_config": {
            "dp_size": 2,
            "cp_size": 1,
            "tp_size": 1,
            "sp_size": 1,
            "ep_size": 1,
            "pp_size": 1,
            "world_size": 2  # 需要2个设备
        },
        "ranktable": {
            "version": "1.0",
            "status": "active",
            "server_count": "1",
            "server_list": [
                {
                    "server_id": "server_0",
                    "host_ip": "192.168.1.100",
                    "device": [
                        {
                            "device_id": "0",
                            "device_ip": "192.168.1.100",
                            "super_device_id": None,
                            "rank_id": "0"
                        },
                        {
                            "device_id": "1",
                            "device_ip": "192.168.1.100",
                            "super_device_id": None,
                            "rank_id": "1"
                        }
                    ]
                }
            ]
        }
    }
    
    try:
        print("发送多设备注册请求...")
        print(f"数据: {json.dumps(register_data, indent=2, ensure_ascii=False)}")
        
        response = requests.post(controller_url, json=register_data, timeout=30)
        
        if response.status_code == 200:
            result = response.json()
            print(f"多设备注册成功! 响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
        else:
            print(f"多设备注册失败! 状态码: {response.status_code}")
            print(f"错误信息: {response.text}")
            
    except Exception as e:
        print(f"多设备注册发生错误: {e}")


if __name__ == "__main__":
    print("=== 单设备注册示例 ===")
    simple_register_example()

    time.sleep(5) # wait for the controller to start
    
    print("\n=== 多设备注册示例 ===")
    multi_device_register_example()
